package utils

import (
	"context"
	"sync"
)

type WorkerPool struct {
	workerCount int
	tasks       chan func() error
	errors      []error
	mu          sync.Mutex
}

func NewWorkerPool(workerCount int) *WorkerPool {
	return &WorkerPool{
		workerCount: workerCount,
		tasks:       make(chan func() error),
		errors:      []error{},
	}
}

func (wp *WorkerPool) Run(ctx context.Context) {
	var wg sync.WaitGroup

	// 启动worker
	for i := 0; i < wp.workerCount; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				case task, ok := <-wp.tasks:
					if !ok {
						return
					}
					if err := task(); err != nil {
						wp.mu.Lock()
						wp.errors = append(wp.errors, err)
						wp.mu.Unlock()
					}
				}
			}
		}()
	}

	wg.Wait()
}

func (wp *WorkerPool) AddTask(task func() error) {
	wp.tasks <- task
}

func (wp *WorkerPool) Close() {
	close(wp.tasks)
}

func (wp *WorkerPool) Errors() []error {
	wp.mu.Lock()
	defer wp.mu.Unlock()
	return wp.errors
}
